import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import confusion_matrix
plt.rcParams['font.sans-serif'] = ['SimHei'] # 显示中文
model = LSTM(vocab_size=vocab_size, 
                hidden_dim=hidden_dim, 
                num_layers=num_layers,
               embedding_dim=embedding_dim, 
               output_dim=output_dim) 
model.load_state_dict(torch.load("./best_model.pt"))
x_input = x_test.long().transpose(1, 0).contiguous()
x_input = x_input.to(device)
output_ = model(x_input)  # torch.Size([32, 2])
y_pred = torch.argmax(output_,dim=-1, keepdim=True)  
# y_pred:[6000,1]
# 生成混淆矩阵
con_mat = confusion_matrix(y_test, y_pred)
sns.heatmap(con_mat, annot=True, fmt='d',cmap = 'Blues',
               xticklabels=[0,1], yticklabels=[0,1]) 
plt.ylabel('实际结果', fontsize=18)
plt.xlabel('预测结果', fontsize=18)
plt.show()  
